import cv2
import numpy as np

def extract_grid(frame):
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150)
    a,contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return contours, edges

def find_largest_contour(contours, min_area=15000):
    if contours:
        contours = [c for c in contours if cv2.contourArea(c) > min_area]
        if contours:
            return max(contours, key=cv2.contourArea)
    return None

def draw_grid_contour(frame, contour):
    cv2.drawContours(frame, [contour], -1, (0, 0, 255), 2)

def draw_grid_cells(frame, grid_x, grid_y, grid_w, grid_h, rows, cols):
    cell_width = grid_w // cols
    cell_height = grid_h // rows
    for row in range(rows):
        for col in range(cols):
            start_x = grid_x + col * cell_width
            start_y = grid_y + row * cell_height
            end_x = start_x + cell_width
            end_y = start_y + cell_height
            cv2.rectangle(frame, (start_x, start_y), (end_x, end_y), (0, 255, 255), 1)

def compute_color_stats(image, mask):
    mean = cv2.mean(image, mask=mask)[:3]
    stddev = cv2.meanStdDev(image, mask=mask)[1].flatten()
    return mean, stddev

first_frame = None
second_frame = None

def main_function(frame,key):
    global first_frame
    global second_frame

    Before = (-1, -1)
    after = (-1, -1)

    contours, edges = extract_grid(frame)

    if key == ord('1'):
        first_frame = frame.copy()
        print("First frame recorded.")

    elif key == ord('2'):
        second_frame = frame.copy()
        print("Second frame recorded.")

    if first_frame is not None and second_frame is not None:
        largest_contour = find_largest_contour(contours, min_area=15000)
        if largest_contour is not None:
            x, y, w, h = cv2.boundingRect(largest_contour)
            draw_grid_contour(frame, largest_contour)
            draw_grid_cells(frame, x, y, w, h, rows=3, cols=3)

            cell_width = w // 3
            cell_height = h // 3

            max_variance_diff = -float('inf')
            Before = (-1, -1)
            min_variance_diff = float('inf')
            after = (-1, -1)

            for row in range(3):
                for col in range(3):
                    start_x = x + col * cell_width
                    start_y = y + row * cell_height
                    end_x = start_x + cell_width
                    end_y = start_y + cell_height

                    cell_frame_first = first_frame[start_y:end_y, start_x:end_x]
                    cell_frame_second = second_frame[start_y:end_y, start_x:end_x]

                    mask = np.ones(cell_frame_first.shape[:2], dtype=np.uint8)

                    _, stddev_first = compute_color_stats(cell_frame_first, mask)
                    _, stddev_second = compute_color_stats(cell_frame_second, mask)

                    # 计算标准差的差值
                    stddev_diff = stddev_first - stddev_second
                    variance_diff = stddev_diff.mean()  # 计算平均差值，可以保留负值

                    if variance_diff > max_variance_diff:
                        max_variance_diff = variance_diff
                        Before = (start_x + cell_width // 2, start_y + cell_height // 2)

                    if variance_diff < min_variance_diff:
                        min_variance_diff = variance_diff
                        after = (start_x + cell_width // 2, start_y + cell_height // 2)

            if Before != (-1, -1):
                print(f"Cell with maximum variance difference is at: {Before}")
                cv2.circle(frame, Before, 10, (0, 255, 0), 2)
                cv2.putText(frame, 'Before', Before, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

            if after != (-1, -1):
                print(f"Cell with minimum variance difference is at: {after}")
                if min_variance_diff < 0:
                    print(f"Minimum variance difference is negative: {min_variance_diff}")
                cv2.circle(frame, after, 10, (255, 0, 0), 2)
                cv2.putText(frame, 'After', after, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)

    if edges is not None:
        pass
        # cv2.imshow('Edge Detection', edges)

    return frame, Before,after
